import tkinter as tk
from tkinter import ttk, filedialog, messagebox
import cv2
import numpy as np
import os
from PIL import Image, ImageTk
import threading
from tqdm import tqdm

class ImageCropper:
    def __init__(self):
        self.crop_width = None  # 裁剪宽度
        self.reference_right_edge = None  # 参考右边缘位置
        self.crop_offset = 100  # 默认向左裁剪的距离
        
    def find_right_edge(self, binary_img):
        # 计算垂直投影
        projection = np.sum(binary_img, axis=0)
        
        # 使用一维高斯卷积减少噪声影响
        window = 15
        sigma = window/3.0
        x = np.linspace(-window//2, window//2, window)
        gaussian = np.exp(-(x**2)/(2*sigma**2))
        gaussian = gaussian/np.sum(gaussian)
        smoothed = np.convolve(projection, gaussian, mode='valid')
        
        # 计算梯度
        gradient = np.gradient(smoothed)
        
        # 从右向左扫描找到第一个显著边缘
        threshold = np.std(gradient) * 2
        for i in range(len(gradient)-1, window, -1):
            if abs(gradient[i]) > threshold:
                local_region = gradient[max(0, i-5):min(len(gradient), i+6)]
                local_offset = np.argmax(np.abs(local_region))
                return i - 5 + local_offset
        return None
    
    def initialize_with_reference(self, reference_img):
        # 转换为灰度图
        gray = cv2.cvtColor(reference_img, cv2.COLOR_BGR2GRAY)
        
        # 增强对比度
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
        gray = clahe.apply(gray)
        
        # 检测垂直线
        binary = detect_vertical_lines(gray)
        
        # 找到右边缘
        right_edge = self.find_right_edge(binary)
        if right_edge is not None:
            self.reference_right_edge = right_edge
            # 设置裁剪宽度（从右边缘向左的距离）
            self.crop_width = right_edge - self.crop_offset  # 使用可调节的偏移量
    
    def crop_image(self, img):
        if self.crop_width is None:
            return img
            
        # 转换为灰度图并增强对比度
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
        gray = clahe.apply(gray)
        
        # 检测垂直线
        binary = detect_vertical_lines(gray)
        
        # 找到当前帧的右边缘
        current_right_edge = self.find_right_edge(binary)
        
        if current_right_edge is None:
            return img
            
        # 计算裁剪区域
        start_x = max(0, current_right_edge - self.crop_width)
        end_x = current_right_edge
        
        # 裁剪图像
        cropped = img[:, start_x:end_x]
        
        return cropped

def detect_vertical_lines(gray_img):
    # 使用Sobel算子增强垂直边缘检测
    sobelx = cv2.Sobel(gray_img, cv2.CV_64F, 1, 0, ksize=5)
    sobelx = cv2.convertScaleAbs(sobelx)
    
    # 使用自适应阈值
    binary = cv2.adaptiveThreshold(sobelx, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                 cv2.THRESH_BINARY, 21, 2)
    
    # 形态学操作增强垂直线条
    kernel_vertical = np.ones((15,1), np.uint8)
    binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel_vertical)
    binary = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel_vertical)
    
    return binary

class PIVProcessorGUI:
    def __init__(self, root):
        self.root = root
        self.root.title("PIV图像处理工具")
        self.root.geometry("1000x800")
        
        # 创建主框架
        self.main_frame = ttk.Frame(self.root, padding="10")
        self.main_frame.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        
        # 存储输入输出路径对
        self.folder_pairs = []  # 存储(输入路径, 输出路径)的元组列表
        
        # 创建图像处理器
        self.cropper = ImageCropper()
        
        self._create_widgets()
        
        # 处理状态
        self.processing = False
        self.preview_image = None
        
        # 显示欢迎信息
        self._show_welcome_message()
        
    def _create_widgets(self):
        # 标题和帮助按钮
        title_frame = ttk.Frame(self.main_frame)
        title_frame.grid(row=0, column=0, columnspan=3, pady=(0, 10))
        ttk.Label(title_frame, text="PIV图像处理工具", font=('Arial', 14, 'bold')).pack(side=tk.LEFT, padx=5)
        ttk.Button(title_frame, text="使用说明", command=self._show_help).pack(side=tk.RIGHT, padx=5)
        
        # 裁剪控制框架
        crop_control_frame = ttk.LabelFrame(self.main_frame, text="裁剪控制", padding="5")
        crop_control_frame.grid(row=1, column=0, columnspan=3, sticky=(tk.W, tk.E), pady=(0, 10))
        
        # 添加裁剪距离滑块
        ttk.Label(crop_control_frame, text="向左裁剪距离:").grid(row=0, column=0, padx=5)
        self.crop_slider = ttk.Scale(crop_control_frame, from_=50, to=300, orient=tk.HORIZONTAL, length=200)
        self.crop_slider.set(100)  # 设置默认值
        self.crop_slider.grid(row=0, column=1, padx=5)
        self.crop_value_label = ttk.Label(crop_control_frame, text="100")
        self.crop_value_label.grid(row=0, column=2, padx=5)
        
        # 绑定滑块值变化事件
        self.crop_slider.configure(command=self._update_crop_value)
        
        # 文件夹列表框架
        folders_frame = ttk.LabelFrame(self.main_frame, text="文件夹对应关系", padding="5")
        folders_frame.grid(row=2, column=0, columnspan=3, sticky=(tk.W, tk.E))
        
        # 创建文件夹列表
        self.folders_list = ttk.Treeview(folders_frame, columns=("输入文件夹", "输出文件夹"), show="headings", height=6)
        self.folders_list.heading("输入文件夹", text="输入文件夹")
        self.folders_list.heading("输出文件夹", text="输出文件夹")
        self.folders_list.grid(row=0, column=0, columnspan=3, sticky=(tk.W, tk.E), padx=5, pady=5)
        
        # 添加滚动条
        scrollbar = ttk.Scrollbar(folders_frame, orient=tk.VERTICAL, command=self.folders_list.yview)
        scrollbar.grid(row=0, column=3, sticky=(tk.N, tk.S))
        self.folders_list.configure(yscrollcommand=scrollbar.set)
        
        # 文件夹操作按钮
        btn_frame = ttk.Frame(folders_frame)
        btn_frame.grid(row=1, column=0, columnspan=3, pady=5)
        ttk.Button(btn_frame, text="添加文件夹对", command=self._add_folder_pair).pack(side=tk.LEFT, padx=5)
        ttk.Button(btn_frame, text="删除选中", command=self._remove_selected_pair).pack(side=tk.LEFT, padx=5)
        ttk.Button(btn_frame, text="清空列表", command=self._clear_folder_pairs).pack(side=tk.LEFT, padx=5)
        
        # 预览区域
        preview_frame = ttk.LabelFrame(self.main_frame, text="预览", padding="5")
        preview_frame.grid(row=3, column=0, columnspan=3, pady=10, sticky=(tk.W, tk.E))
        ttk.Label(preview_frame, text="原始图像", font=('Arial', 10)).grid(row=0, column=0, padx=5)
        ttk.Label(preview_frame, text="处理后图像", font=('Arial', 10)).grid(row=0, column=1, padx=5)
        self.preview_label = ttk.Label(preview_frame)
        self.preview_label.grid(row=1, column=0, columnspan=2, pady=5)
        
        # 控制按钮
        control_frame = ttk.Frame(self.main_frame)
        control_frame.grid(row=4, column=0, columnspan=3, pady=10)
        ttk.Button(control_frame, text="预览", command=self._preview).grid(row=0, column=0, padx=5)
        ttk.Button(control_frame, text="开始处理", command=self._start_processing).grid(row=0, column=1, padx=5)
        
        # 进度条和进度显示框架
        progress_frame = ttk.Frame(self.main_frame)
        progress_frame.grid(row=5, column=0, columnspan=3, pady=10)
        
        # 进度条
        self.progress = ttk.Progressbar(progress_frame, length=300, mode='determinate')
        self.progress.pack(side=tk.LEFT, padx=5)
        
        # 进度数值标签
        self.progress_label = ttk.Label(progress_frame, text="0%")
        self.progress_label.pack(side=tk.LEFT, padx=5)
        
    def _add_folder_pair(self):
        input_path = filedialog.askdirectory(title="选择输入文件夹")
        if not input_path:
            return
            
        output_path = filedialog.askdirectory(title="选择对应的输出文件夹")
        if not output_path:
            return
            
        # 添加到列表
        self.folder_pairs.append((input_path, output_path))
        self.folders_list.insert("", "end", values=(input_path, output_path))
        
    def _remove_selected_pair(self):
        selected_item = self.folders_list.selection()
        if not selected_item:
            return
            
        # 获取选中项的索引
        index = self.folders_list.index(selected_item)
        # 从数据和显示中删除
        self.folder_pairs.pop(index)
        self.folders_list.delete(selected_item)
        
    def _clear_folder_pairs(self):
        self.folder_pairs.clear()
        self.folders_list.delete(*self.folders_list.get_children())
        
    def _update_crop_value(self, value):
        # 更新裁剪值显示
        self.crop_value_label.config(text=str(int(float(value))))
        self.cropper.crop_offset = int(float(value))
        # 如果已经初始化过裁剪器，更新裁剪宽度
        if self.cropper.reference_right_edge is not None:
            self.cropper.crop_width = self.cropper.reference_right_edge - self.cropper.crop_offset
        
    def _preview(self):
        if not self.folder_pairs:
            messagebox.showerror("错误", "请至少添加一对文件夹")
            return
            
        # 显示文件夹对应关系确认对话框
        self._show_folder_pairs_confirmation()
        
        # 预览第一对文件夹的第一张图片
        input_path = self.folder_pairs[0][0]
        
        # 获取第一张图片
        image_files = [f for f in os.listdir(input_path) if f.lower().endswith('.bmp')]
        if not image_files:
            messagebox.showerror("错误", f"在文件夹 {input_path} 中没有找到BMP文件")
            return
            
        # 读取第一张图片
        img_path = os.path.join(input_path, image_files[0])
        img = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), cv2.IMREAD_COLOR)
        
        # 初始化裁剪器
        self.cropper.initialize_with_reference(img)
        
        # 处理图片
        processed = self.cropper.crop_image(img)
        
        # 显示预览
        self._show_preview(img, processed)
        
    def _show_folder_pairs_confirmation(self):
        confirm_text = "请确认以下文件夹对应关系：\n\n"
        for i, (input_path, output_path) in enumerate(self.folder_pairs, 1):
            confirm_text += f"{i}. 输入：{input_path}\n   输出：{output_path}\n\n"
        
        dialog = tk.Toplevel(self.root)
        dialog.title("确认文件夹对应关系")
        dialog.geometry("600x400")
        
        # 使用文本框显示对应关系
        text_widget = tk.Text(dialog, wrap=tk.WORD, padx=10, pady=10)
        text_widget.pack(fill=tk.BOTH, expand=True)
        text_widget.insert(tk.END, confirm_text)
        text_widget.config(state=tk.DISABLED)
        
        # 添加滚动条
        scrollbar = ttk.Scrollbar(dialog, orient=tk.VERTICAL, command=text_widget.yview)
        scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
        text_widget.config(yscrollcommand=scrollbar.set)
        
        # 确认按钮
        ttk.Button(dialog, text="确认", command=dialog.destroy).pack(pady=10)
        
    def _process_images(self):
        try:
            total_progress = 0
            total_files = 0
            
            # 计算总文件数
            for input_path, _ in self.folder_pairs:
                image_files = [f for f in os.listdir(input_path) if f.lower().endswith('.bmp')]
                total_files += len(image_files)
            
            # 处理每对文件夹
            for input_path, output_path in self.folder_pairs:
                # 创建输出目录
                os.makedirs(output_path, exist_ok=True)
                
                # 获取所有图片
                image_files = [f for f in os.listdir(input_path) if f.lower().endswith('.bmp')]
                
                # 读取第一张图片并初始化裁剪器
                if image_files:
                    first_img_path = os.path.join(input_path, image_files[0])
                    first_img = cv2.imdecode(np.fromfile(first_img_path, dtype=np.uint8), cv2.IMREAD_COLOR)
                    self.cropper.initialize_with_reference(first_img)
                
                # 处理每张图片
                for img_file in image_files:
                    # 更新进度条和进度标签
                    total_progress += 1
                    progress_value = (total_progress / total_files) * 100
                    self.progress['value'] = progress_value
                    self.progress_label.config(text=f"{progress_value:.1f}%")
                    self.root.update_idletasks()
                    
                    # 处理图片
                    img_path = os.path.join(input_path, img_file)
                    img = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), cv2.IMREAD_COLOR)
                    processed = self.cropper.crop_image(img)
                    
                    # 保存处理后的图片
                    output_path_file = os.path.join(output_path, img_file)
                    cv2.imencode('.bmp', processed)[1].tofile(output_path_file)
            
            messagebox.showinfo("完成", "所有文件夹的图片处理完成！")
            
        except Exception as e:
            messagebox.showerror("错误", f"处理过程中发生错误: {str(e)}")
            
        finally:
            self.processing = False
            self.progress['value'] = 0
            self.progress_label.config(text="0%")

    def _show_help(self):
        help_text = """
PIV图像处理工具使用说明

1. 基本操作流程：
   ① 点击"添加文件夹对"按钮选择输入和对应的输出文件夹
   ② 可以添加多对文件夹，按添加顺序一一对应
   ③ 点击"预览"按钮查看效果（会先显示文件夹对应关系确认）
   ④ 确认无误后点击"开始处理"进行批量处理

2. 功能说明：
   - 支持多文件夹批量处理
   - 自动检测图像右侧边缘
   - 根据检测到的边缘进行智能裁剪
   - 对裁剪后的图像进行二值化处理
   - 实时显示处理进度

3. 文件夹管理：
   - "添加文件夹对"：选择一对输入输出文件夹
   - "删除选中"：删除列表中选中的文件夹对
   - "清空列表"：清除所有文件夹对

4. 注意事项：
   - 处理前请确认文件夹对应关系是否正确
   - 输入文件夹必须包含BMP格式图片
   - 处理过程中请勿关闭程序
   - 确保输出文件夹有足够的存储空间

5. 最佳实践：
   - 先用少量图片测试效果
   - 确认效果满意后再处理大量图片
   - 定期备份重要数据
        """
        
        help_dialog = tk.Toplevel(self.root)
        help_dialog.title("使用说明")
        help_dialog.geometry("600x500")
        
        # 使用文本框显示帮助信息
        text_widget = tk.Text(help_dialog, wrap=tk.WORD, padx=10, pady=10)
        text_widget.pack(fill=tk.BOTH, expand=True)
        text_widget.insert(tk.END, help_text)
        text_widget.config(state=tk.DISABLED)
        
        # 添加滚动条
        scrollbar = ttk.Scrollbar(help_dialog, orient=tk.VERTICAL, command=text_widget.yview)
        scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
        text_widget.config(yscrollcommand=scrollbar.set)
        
    def _show_welcome_message(self):
        welcome_text = """欢迎使用PIV图像处理工具！

本工具支持多文件夹批量处理功能：
1. 可以同时添加多对输入输出文件夹
2. 文件夹按添加顺序一一对应
3. 处理前会显示对应关系确认

点击右上角的"使用说明"按钮可以查看详细的使用指南。

祝您使用愉快！"""
        
        messagebox.showinfo("欢迎", welcome_text)

    def _start_processing(self):
        if self.processing:
            return
            
        if not self.folder_pairs:
            messagebox.showerror("错误", "请至少添加一对文件夹")
            return
            
        # 在新线程中启动处理
        self.processing = True
        thread = threading.Thread(target=self._process_images)
        thread.start()
        
    def _show_preview(self, original, processed):
        # 调整图像大小以适应显示
        height = 200
        aspect_ratio = original.shape[1] / original.shape[0]
        width = int(height * aspect_ratio)
        
        # 调整大小
        original_resized = cv2.resize(original, (width, height))
        processed_resized = cv2.resize(processed, (width, height))
        
        # 创建间隔
        separator = np.ones((height, 20, 3), dtype=np.uint8) * 128  # 灰色间隔
        
        # 水平拼接（添加间隔）
        combined = np.hstack((original_resized, separator, processed_resized))
        
        # 转换为PIL格式
        combined = cv2.cvtColor(combined, cv2.COLOR_BGR2RGB)
        img = Image.fromarray(combined)
        img_tk = ImageTk.PhotoImage(image=img)
        
        # 更新预览
        self.preview_label.configure(image=img_tk)
        self.preview_label.image = img_tk

if __name__ == '__main__':
    root = tk.Tk()
    app = PIVProcessorGUI(root)
    root.mainloop() 